package com.suxin.ratelimit.filters;

import com.suxin.ratelimit.Rate;
import com.suxin.ratelimit.properties.RateLimiterProperties;
import com.suxin.ratelimit.repository.RateLimiter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;
import java.util.List;
import java.util.StringJoiner;

/**
 * @author Tang
 * @classname RateLimiterWrapperFilter
 * @description [ 自定义一个filter,实现限流 ]
 * @date 2021/7/14 23:41
 */
public class RateLimiterWrapperFilter implements GlobalFilter {

    @Resource
    RateLimiterProperties rateLimiterProperties;

    @Resource
    RateLimiter rateLimiter;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        if (!rateLimiterProperties.isEnable()) {
            chain.filter(exchange);
            return Mono.empty();
        }
        RateLimiterProperties.Policy policy = rateLimiterProperties.getDefaultPolicy();
        List<RateLimiterProperties.Policy.Type> policyTypes = policy.getTypes();
        ServerHttpRequest request = exchange.getRequest();
        if (policy != null) {
            String newKey = key(policyTypes,request);
            // 消费限流
            Rate rate = rateLimiter.consumer(policy, newKey);
            if (rate.getLimit() < 0) {
                // todo to many request
            }
        }
        return Mono.empty();
    }

    /**
     * 构建key
     * @param policyTypes
     * @param request
     * @return
     */
    private String key(List<RateLimiterProperties.Policy.Type> policyTypes,ServerHttpRequest request) {
        StringJoiner joiner = new StringJoiner(":");
        policyTypes.forEach(item -> {
            switch (item) {
                case URL:
                    joiner.add(request.getURI().getPath());
                    break;
                case ORIGIN:
                    joiner.add(request.getLocalAddress().getAddress().toString());
                    break;
                case USER:
                default:
                    break;
            }
        });
        return joiner.toString();
    }

}